import os
import requests
import shutil
import time
import pdf_parser
import paper_reviewer
import configs
from datetime import datetime
from flask import Flask, render_template, request, Response, send_file, jsonify
from flask_cors import CORS

app = Flask(__name__, template_folder='templates')
CORS(app)

base_conf = configs.Base()
nougat_conf = configs.Nougat()

PDF_DIR = base_conf.data_dir
if not os.path.exists(PDF_DIR):
    os.makedirs(PDF_DIR)

MMD_DIR = base_conf.data_dir
if not os.path.exists(MMD_DIR):
    os.makedirs(MMD_DIR)

PAGE_DIR = base_conf.data_dir
if not os.path.exists(PAGE_DIR):
    os.makedirs(PAGE_DIR)

def update_filename(filename, timestamp):
    #filename = filename.replace(' ', '_')
    base, ext = os.path.splitext(filename)
    return f'{base}-{timestamp}{ext}'

def get_pdf_path(file_id):
    return f'{PDF_DIR}/{file_id}.pdf'

def get_mmd_path(file_id):
    return f'{MMD_DIR}/{file_id}.mmd'

def get_page_path(file_id):
    return f'{PAGE_DIR}/{file_id}_page.png'

@app.route(f'/page/<path:file_id>')
def serve_page(file_id):
    page_path = get_page_path(file_id)
    print(f'serve_page: {page_path}')
    return send_file(page_path, mimetype='application/png')

@app.route(f'/pdf/<path:file_id>')
def serve_pdf(file_id):
    pdf_path = get_pdf_path(file_id)
    print(f'serve_pdf: {pdf_path}')
    return send_file(pdf_path, mimetype='application/pdf')

@app.route('/', methods=['GET', 'POST'])
def index():
    user_agent = request.user_agent.string
    is_mobile = 'Android' in user_agent or 'webOS' in user_agent or 'iPhone' in user_agent or \
                'iPad' in user_agent or 'iPod' in user_agent or 'BlackBerry' in user_agent or \
                'IEMobile' in user_agent or 'Opera Mini' in user_agent
    device_type = 'mobile' if is_mobile else 'desktop'

    if request.method == 'GET':
        if device_type == 'mobile':
            return render_template('index_mobile.html', device_type=device_type)
        else:
            return render_template('index.html', device_type=device_type)
    else:
        file_id = request.form.get('file_id')
        print(f'file_id={file_id}')
        if device_type =='mobile':
            return render_template('index_mobile.html', device_type=device_type, file_id=file_id)
        else:
            return render_template('index.html', device_type=device_type, file_id=file_id)

@app.route('/upload', methods=['POST'])
def parse():
    timestamp = datetime.now().strftime("%y%m%d%H%M%S")
    print(f'timestamp={timestamp}')
    file = request.files.get('file')
    url = request.form.get('url')
    if file:
        filename = update_filename(file.filename, timestamp)
        pdf_path = os.path.join(PDF_DIR, filename)
        file.save(pdf_path)
    elif url:
        try:
            os.environ['https_proxy'] = base_conf.https_proxy
            response = requests.get(url, stream=True)
            filename = url.split("/")[-1]
            if not filename.endswith('.pdf'):
                filename += '.pdf'
            filename = update_filename(filename, timestamp)
            pdf_path = os.path.join(PDF_DIR, filename)
            with open(pdf_path, 'wb') as out_file:
                shutil.copyfileobj(response.raw, out_file)
        except Exception as e:
            return jsonify(dict(status='error', message=f'Error downloading file: {e}')), 500
        finally:
            os.environ['https_proxy'] = ''
    else:
        return jsonify(dict(status='error', message='No file or URL provided')), 400

    stop_page = None
    if hasattr(nougat_conf, 'stop_page') and nougat_conf.stop_page > 0:
        stop_page = nougat_conf.stop_page

    mmd_path, page_path = pdf_parser.run_parse(pdf_path, MMD_DIR, stop=stop_page)

    file_id, _ = os.path.splitext(filename)
    if mmd_path or page_path:
        print(f'Succeeded - file_id={file_id}')
        return jsonify(dict(status='success', file_id=file_id))
    else:
        print(f'Failed - file_id={file_id}')
        return jsonify(dict(status='error', message=f'Failed to parse file: {file_id}')), 500

@app.route('/review', methods=['POST'])
def review():
    file_id = request.form.get('file_id')
    mmd_path = get_mmd_path(file_id)
    def generate():
        ans = paper_reviewer.run_review(mmd_path)
        buff = ''
        for chunk in ans:
            if chunk and chunk.choices:
                msg = chunk.choices[0].delta
                if msg.content:
                    buff += msg.content
                    if len(buff) > 20:
                        yield buff
                        buff = ''
    return Response(generate(), mimetype='text/markdown')

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=base_conf.app_port)

